import multiprocessing
import shutil
from multiprocessing import Pool

from batchgenerators.utilities.file_and_folder_operations import *

from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json
from nnunetv2.paths import nnUNet_raw
from skimage import io,color,img_as_ubyte
from acvl_utils.morphology.morphology_helper import generic_filter_components
from scipy.ndimage import binary_fill_holes


def load_and_covnert_case(input_image: str, input_seg: str, output_image: str, output_seg: str,
                          min_component_size: int = 50):
    seg = io.imread(input_seg)
    seg = color.rgb2gray(seg)
    seg = img_as_ubyte(seg)
    seg[seg >0] = 1
    # image = io.imread(input_image)
    # image = image.sum(2)
    # mask = image == (3 * 255)
    # # the dataset has large white areas in which road segmentations can exist but no image information is available.
    # # Remove the road label in these areas
    # mask = generic_filter_components(mask, filter_fn=lambda ids, sizes: [i for j, i in enumerate(ids) if
    #                                                                      sizes[j] > min_component_size])
    # mask = binary_fill_holes(mask)
    # seg[mask] = 0
    io.imsave(output_seg, seg, check_contrast=False)
    shutil.copy(input_image, output_image)


if __name__ == "__main__":

    # extracted archive from https://www.kaggle.com/datasets/insaff/massachusetts-roads-dataset?resource=download
    source = '/media/alex/USB DISK/images/image'
    

    dataset_name = 'Dataset890_CellsegZigong'

    imagestr = join(nnUNet_raw, dataset_name, 'imagesTr')
    imagests = join(nnUNet_raw, dataset_name, 'imagesTs')
    labelstr = join(nnUNet_raw, dataset_name, 'labelsTr')
    labelsts = join(nnUNet_raw, dataset_name, 'labelsTs')
    maybe_mkdir_p(imagestr)
    maybe_mkdir_p(imagests)
    maybe_mkdir_p(labelstr)
    maybe_mkdir_p(labelsts)

    train_source = join(source, 'cellTrain')
    maybe_mkdir_p(train_source)
    maybe_mkdir_p(join(train_source, 'image'))
    maybe_mkdir_p(join(train_source, 'mask'))

    for file in subfiles(source, join=False, suffix='_mask.png'):
        shutil.copy(join(source,file),join(train_source,'mask',file.replace('_mask','')))
        shutil.copy(join(source,file.replace('_mask','')),join(train_source,'image',file.replace('_mask','')))


    valid_ids = subfiles(join(train_source, 'mask'), join=False, suffix='png')
    num_train = len(valid_ids)
    for v in valid_ids:
        load_and_covnert_case(
                join(train_source, 'image', v),
                join(train_source, 'mask', v),
                join(imagestr, v[:-4] + '_0000.png'),
                join(labelstr, v),
                50)
       
    generate_dataset_json(join(nnUNet_raw, dataset_name), {0: 'R', 1: 'G', 2: 'B'}, {'background': 0, 'cell':1},
                          num_train, '.png', dataset_name=dataset_name)
